#include "precompiled.h"
#include "IMLGenerator.h"

namespace IcyCrystal {
	Symbol* IMLGenerator::createTempSymbol(Namespace* currentNamespace, Arg& arg){
		// Is it a temporary? Add it to symbol map.
		if( arg.arg == TEMP && currentNamespace->SymbolMap.count(*arg.name.symbolName) == 0){
			Symbol* newTemp = new Symbol;
			newTemp->scope = SCOPE_TEMP;
			newTemp->level = currentNamespace->level;
			newTemp->symbolName = *arg.name.symbolName;
			//delete arg.name.symbolName;
			arg.name.symbolName = & newTemp->symbolName;
			newTemp->type.primitiveType  = PRIM_TEMP;
			newTemp->location = new Arg;
			newTemp->location= &arg;
			newTemp->locations = new LocationList;
			newTemp->locations->push_back(newTemp->location);

			currentNamespace->SymbolMap[newTemp->symbolName] = newTemp;

			return newTemp;
		}
		else if( arg.arg == TEMP || arg.arg == LOCAL){
			return currentNamespace->findSymbol(*arg.name.symbolName);
		}
		else if(arg.arg == GLOBAL){
			currentNamespace->SymbolMap[*arg.name.symbolName] = globalNamespace.SymbolMap[*arg.name.symbolName];
			Arg local;
			local.arg = LOCAL;
			local.name.symbolName = arg.name.symbolName;
			return createTempSymbol(currentNamespace, local);
		}
		else if( arg.arg == CONSTANT ){
			return globalNamespace.ConstantMap[*arg.name.symbolName];
		}

		return NULL;
	}
	

	void IMLGenerator::computeLastUseHelper(Namespace* currentNamespace, Arg& arg, UseMap& map, int i, int x){
		if( arg.arg == LOCAL || arg.arg == GLOBAL || arg.arg == TEMP){
			// Check if exists.
			if( map.count(arg.name.symbolName) == 0 ){
				// Add it.
				UseInfo info;
				info.lastUse = i - x;
				map[arg.name.symbolName] = info;
			}
		}		

	}

	// Load a symbol into a register
	Quad* IMLGenerator::generateLoadQuad(Symbol* sym, int reg){
		Quad* ret = new Quad;
		ret->op = MOV;
		ret->inner.normalOp.arg1.arg = scopeToArgType(sym);
		ret->inner.normalOp.arg1.name.symbolName = &sym->symbolName;

		ret->inner.normalOp.arg2.arg = REGISTER;
		ret->inner.normalOp.arg2.name.intName = reg;

		// Change R so it holds *only* x
		registerArray[reg].variables->clear();
		registerArray[reg].variables->push_back(sym);

		// Add R as a location for x
		sym->locations->push_back(&ret->inner.normalOp.arg2);

		return ret;
	}

	// Restore a symbol from a register.
	Quad* IMLGenerator::generateStoreQuad(Symbol* sym, int reg){
		Quad* ret = new Quad;
		ret->op = MOV;
		ret->inner.normalOp.arg1.arg = REGISTER;
		ret->inner.normalOp.arg1.name.intName = reg;

		ret->inner.normalOp.arg2.arg = scopeToArgType(sym);
		ret->inner.normalOp.arg2.name.symbolName = &sym->symbolName;

		// Change x to include its own memory location (if not already)
		bool found =false;
		LocationList::iterator it;
		for( it = sym->locations->begin(); it != sym->locations->end(); it++){
			Arg* arg = *it;
			if( arg == sym->location ){
				found = true;
				break;
			}
		}

		if(!found){
			sym->locations->push_back(sym->location);
		}

		return ret;
	}

	UseMap* IMLGenerator::computeLastUse(Namespace* currentNamespace, ArgList* block){
		UseMap* ret = new UseMap;
		UseMap& map = *ret;

		int i = block->size();
		int x = 0;

		// Walk the block backwards.
		ArgList::reverse_iterator it;
		for( it = block->rbegin(); it!= block->rend(); it++){
			Quad* quad = *it;
			// Check the operands.
			switch(quad->op){
			case ADD:
			case SUB:
			case MULT:
			case DIV:
				// Two operands.
				computeLastUseHelper(currentNamespace, quad->inner.normalOp.arg1,map,i,x);
				computeLastUseHelper(currentNamespace, quad->inner.normalOp.arg2,map,i,x);
				break;
			case PUSH:
			case POP:
			case MOV:
				computeLastUseHelper(currentNamespace, quad->inner.normalOp.arg1,map,i,x);
				break;
			}

			x++;
		}

		return ret;
	}

	ArgList* IMLGenerator::generateEndBlock(Namespace* currentNamespace){
		ArgList* ret = new ArgList;
		// Walk through symbols.
		_symbolMap::iterator it;
		for(it = currentNamespace->SymbolMap.begin(); it != currentNamespace->SymbolMap.end(); it++){
			Symbol* sym = it->second;

			if( sym->scope == SCOPE_GLOBAL || sym->scope == SCOPE_LOCAL || sym->scope == SCOPE_PARAM){
				// Look for its locations.
				LocationList::iterator i;
				bool found = false;
				Arg foundRegister = Arg();
				for( i = sym->locations->begin(); i != sym->locations->end(); i++){
					Arg* loc = *i;
					if( loc->arg == LOCAL || loc->arg == GLOBAL){
						if( loc->name.symbolName->compare(sym->symbolName) == 0 ){
							found = true;
						}
					}
					else if( loc->arg == STACK ){
						found = true;
					}
					else if( loc->arg == REGISTER ){
						foundRegister = *loc;
					}
				}

				if( ! found ){
					// Generate a simple store.
					Quad* simStore = new Quad;
					simStore->op = MOV;
					simStore->inner.normalOp.arg1.arg = REGISTER;
					simStore->inner.normalOp.arg1.name.intName = foundRegister.name.intName;

					simStore->inner.normalOp.arg2.arg = scopeToArgType(sym);
					simStore->inner.normalOp.arg2.name.symbolName = & sym->symbolName;

					ret->push_back(simStore);
				}
			}


			// Clear the symbol's location.
			sym->locations->clear();
			sym->locations->push_back(sym->location);
		}

		// Clear all registers.
		clearAllRegisters();
		return ret;
	}

	ArgList* IMLGenerator::doSecondPass( Namespace* currentNamespace, ArgList* block){
		ArgList* ret = new ArgList;
		// Create the lastUse map to help the register functions.
		UseMap* map = computeLastUse(currentNamespace, block);

		// Go through each statement.
		ArgList::iterator it;
		int x = 0;
		for(it=block->begin();it!=block->end();it++){
			Quad* quad = *it;
			bool isTwoOps = true;
			bool targetIsZ = false;
			bool isMov = false;
			switch (quad->op){
			case ADD:
			case SUB:
			case MULT:
			case DIV:
				isTwoOps = true;
				targetIsZ = true;
				break;
			case PUSH:
			case POP:
				isTwoOps = false;
			case RETURN:
				continue;
			case MOV:
				isMov = true;
			}
			if( ! isMov ) {
				// Get Rx, Ry, Rz
				// x = y + z
				// Find the symbols for each variable.
				Symbol* sX, *sY, *sZ;
				ArgList* rxIns, *ryIns, *rzIns;
				sX = createTempSymbol(currentNamespace, quad->inner.normalOp.result);
				sY = createTempSymbol(currentNamespace, quad->inner.normalOp.arg1);
				if( isTwoOps)
					sZ = createTempSymbol(currentNamespace, quad->inner.normalOp.arg2);
				int Rx = -1, Ry = -1, Rz = -1;
				rxIns = getRegister(sX,*quad,*map,x,Rx);

				// Rx(Rz) should hold ONLY t
				//registerArray[Rx].variables->clear();
				//registerArray[Rx].variables->push_back(sX);

				// Check if y or z are in the registers.
				ryIns = getRegister(sY,*quad,*map,x,Ry);
				if( !symbolIsInRegister(sY, Ry )){
					ret->push_back(generateLoadQuad(sY,Ry));
				}

				if( isTwoOps){
					rzIns = getRegister(sZ,*quad,*map,x,Rz);
					if( isTwoOps && !symbolIsInRegister(sZ,Rz)){
						ret->push_back(generateLoadQuad(sZ,Rz));
					}
				}

				if( rxIns != NULL ){
					ret->insert(ret->end(), rxIns->begin(), rxIns->end());
					delete rxIns;
				}
				if( ryIns != NULL ){
					ret->insert(ret->end(), ryIns->begin(), ryIns->end());
					delete ryIns;
				}
				if( isTwoOps &&rzIns != NULL ){
					ret->insert(ret->end(), rzIns->begin(), rzIns->end());
					delete rzIns;
				}


				// In the case of add and subtract, ( and other similar commands ), z is the target.
				Rx = Rz;

				// Create the add.
				ArgList* genericQuad = generateGenericQuad(currentNamespace, x, *map, quad->op, Ry,Rx, sY, sZ,sX, quad->inner.normalOp.result);
				ret->insert(ret->end(), genericQuad->begin(), genericQuad->end());
				delete genericQuad;
			}
			else {
				// Special case: Copy assignment.
				Symbol *sX, *sY;
				int Ry = -1, Rx = -1;

				bool arg1IsConstant = (quad->inner.normalOp.arg1.arg == INT_CONSTANT || quad->inner.normalOp.arg1.arg == FLOAT_CONSTANT);

				sX = createTempSymbol(currentNamespace, quad->inner.normalOp.result);
				
				if( !arg1IsConstant )
					sY = createTempSymbol(currentNamespace, quad->inner.normalOp.arg1);
				ArgList* rxIns, *ryIns;
				if( arg1IsConstant ){
					// Trick next code since we only need Sx
					sY = sX;
				}
				ryIns = getRegister(sY, *quad, *map, x, Ry);
				if( ryIns != NULL ){
					ret->insert(ret->end(), ryIns->begin(), ryIns->end());
					delete ryIns;
				}
				Rx = Ry;

				if( !arg1IsConstant &&  !symbolIsInRegister(sY, Ry) ){
					ret->push_back(generateLoadQuad(sY,Ry));
				}
				else if( arg1IsConstant ){
					// Generate a simple constant load.
					Quad* simpLoad = new Quad;
					simpLoad->op = MOV;
					// TODO: Handle floats 
					simpLoad->inner.normalOp.arg1.arg = INT_CONSTANT;
					simpLoad->inner.normalOp.arg1.name.intName = quad->inner.normalOp.arg1.name.intName;
					simpLoad->inner.normalOp.arg2.arg = REGISTER;
					simpLoad->inner.normalOp.arg2.name.intName = Rx;
					ret->push_back(simpLoad);
				}

				// Adjust Ry as to include sX
				registerArray[Ry].variables->push_back(sX);
				
				// sX's only location is Ry
				Arg* regArg = new Arg;
				regArg->arg = REGISTER;
				regArg->name.intName = Ry;
				sX->locations->clear();
				sX->locations->push_back(regArg);
				
			}
			x++;
		}

		// Save all symbols that are only in registers to memory.
		ArgList* saveSyms = generateEndBlock(currentNamespace);
		ret->insert(ret->end(), saveSyms->begin(), saveSyms->end());
		return ret;
	}

	ArgList* IMLGenerator::generateGenericQuad(Namespace* currentNamespace, int currentStatement, UseMap& map, OpType op, int lReg, int rReg, Symbol* lSym, Symbol* rSym, Symbol* t, Arg& argT){
		ArgList* retList= new ArgList;
		Quad* ret = new Quad;

		ret->op = op;
		ret->inner.normalOp.arg1.arg = REGISTER;
		ret->inner.normalOp.arg1.name.intName = lReg;

		ret->inner.normalOp.arg2.arg = REGISTER;
		ret->inner.normalOp.arg2.name.intName = rReg;

		// Make sure rReg isnt rSymbol's only location.
		if( rSym->scope != SCOPE_TEMP && rSym->locations->size() == 1 && rSym->locations->at(0)->name.intName == rReg ){
			// Store rSym in its location.
			// Will rSymbol be used again?
			if( map[&rSym->symbolName].lastUse >= currentStatement )
			{
				Quad* simStore = generateStoreQuad(rSym,rReg);
				retList->push_back(simStore);
			}
		}

		registerArray[rReg].variables->clear();
		registerArray[rReg].variables->push_back(t);
		// t's location is now ONLY Rx(Rz)
		t->locations->clear();
		t->locations->push_back(&ret->inner.normalOp.arg2);

		// Remove Rx from the memory location of any other than t
		removeRegisterFromAllSymbols(currentNamespace,rReg, t);

		retList->push_back(ret);

		return retList;
	}

	Quad* IMLGenerator::generateCopyQuad(Namespace* currentNamespace, int rReg, Symbol* lSym){
		Quad* ret = new Quad;

		ret->op = MOV;
		ret->inner.normalOp.arg1.arg = REGISTER;
		ret->inner.normalOp.arg1.name.intName = rReg;

		return ret;
	}
}